#! /usr/bin/env python3
from pyspark.sql.functions import udf

from gai.v2.spark.transformer.column_mapper import ColumnMapper
from gai.v2.spark.transformer.feature_retriever_v2 import target_day


@udf
def _feature_partition_mapping(date: str) -> str:
    return target_day(date)


class DayRectifier(ColumnMapper):
    """A ``DayRectifier`` transforms an input day column to an output day column
    according to certain built-in strategy.

    Args:
        inputCol: the name of the input day column
        outputCol: the name of the output day column

    >>> from gai.v2.utils.platform import get_or_create_spark_session
    >>> from gai.v2.spark.transformer import DayRectifier
    >>> spark = get_or_create_spark_session()
    >>> sample = spark.createDataFrame([('20180304','Alice'), ('20171230', 'Bob')],
    ...                                schema=["actual_day", "name"])
    >>> day_rectifier = DayRectifier(inputCol='actual_day', outputCol='query_day')
    >>> output = day_rectifier.transform(sample)
    >>> output.show()
    +----------+-----+---------+
    |actual_day| name|query_day|
    +----------+-----+---------+
    |  20180304|Alice| 20180226|
    |  20171230|  Bob| 20171225|
    +----------+-----+---------+
    <BLANKLINE>
    >>> output.printSchema()
    root
     |-- actual_day: string (nullable = true)
     |-- name: string (nullable = true)
     |-- query_day: string (nullable = true)
    <BLANKLINE>
    """

    def __init__(self, inputCol, outputCol):
        super(DayRectifier, self).__init__(fun=_feature_partition_mapping,
                                           inputCols=[inputCol],
                                           outputCol=outputCol)
